#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <math.h>
#include <mpi.h>



/* Could use these macros */
#define GET_BIT(var,pos) ((var) & (1<<(pos)))
#define SET_BIT(var,pos) ((var) |= (1<<(pos)))
#define TAG 123

/* 1B prime 22801763489 */
/* 1M prime 15485863 */
/* GRIDSIZE 100000000 takes 15.07 seconds (with k^2>GRIDSIZE optimisation 31.79 without */

/* Globals */
int rank, size; /* MPI */

/* Functions */
int initMPI(int *pargc, char *pargv[]);

int main(int argc, char *argv[])
{
	clock_t c0, c1;
	unsigned short *eratosthenes;
	unsigned long k, l;
	unsigned int process;
	unsigned long start, end, gridpart, gridsize;
	double offset_power;
	long recvValue = 0;
	MPI_Status status;

	c0=clock();

	initMPI(&argc, argv); /* setup MPI, size and rank */

	/* set up range parameters */
	gridsize = atol(argv[1]); /* total size */
	offset_power = atof(argv[2]); /* affects distribution of the grid between processes, 1 is equal, >1 greater ranks have more of the grid */
	
	start = pow(rank, offset_power)*gridsize/pow(size, offset_power)+1;
	end = pow(rank+1, offset_power)*gridsize/pow(size, offset_power);
	gridpart = end - start;
	 
	/* gridpart = gridsize / size; */
	/* changing size of the parts may optimise the algorithmn */
	/* start = (rank == 0) ? 0 : (rank * gridpart) + 1; */
	/* end = (rank == (size - 1)) ? gridsize : (rank + 1) * gridpart; */
	/* gridpart = end - start; */ /* Last rank may have a different size */

	printf("Rank: %d, Gridsize: %ld, Gridpart: %ld, Start: %ld, End: %ld\n", rank, gridsize, gridpart, start, end);

	/* Need to use calloc to initialise values */
	eratosthenes = calloc(ceil(gridpart/8.0), sizeof(unsigned short));

	if (rank == 0)
	{
		eratosthenes[0] |= (1 << (0)); /* Cross off 1 */
		/* eratosthenes[0] |= (1 << (1)); */ /* Cross off 1 */
	}
	else
	{
		while(pow(recvValue,2) < end)
		{
			/* receive and cross off */
			MPI_Recv(&recvValue, 1, MPI_LONG, MPI_ANY_SOURCE, TAG, MPI_COMM_WORLD, &status);
			if (recvValue > 0)
			{
				long offset = recvValue - (start % recvValue);
				/* printf("Rank: %d, Received value: %ld, Offset: %ld\n", rank, recvValue, offset); */
				for(k=offset; k<=gridpart; k+=recvValue)
				{
					/* printf("Rank: %d, (k=%ld) crossing off: %ld\n", rank, k, k+start); */
					eratosthenes[k/8] |= (1 << (k%8));
				}
			}
			else if (recvValue < 0)
			{
				/* printf("Rank: %d, received: %ld with option to break\n", rank, recvValue); */
				if (recvValue >= start)
					break;
			}
		}
	}


	/* TODO: set k correctly, rank 0 should start from 2, others from the base of the grid */
	for(k=0; k<gridpart; k++)
	{
		/* it is prime */
		unsigned long prime = k+start;
		/* printf("Rank: %d testing %ld, flag: %d\n", rank, prime, ((eratosthenes[k/8] & (1 << (k % 8))) >> (k % 8))); */
		/* if it has not been crossed off */
		if (((eratosthenes[k/8] & (1 << (k % 8))) >> (k % 8)) != 1)
		{


			/* Notify other ranks > than this rank */			
			for(process=rank+1; process<size; process++)
			{
				if (process != rank)
				{
					/* printf("Rank: %d sending %ld to process %d\n", rank, prime, process); */
					MPI_Send(&prime, 1, MPI_LONG, process, TAG, MPI_COMM_WORLD);
				}
			}

			/* cross off factors */
			for(l=k+prime; l<=gridpart; l+=prime)
			{
				/* printf("Rank: %d, crossing off: %ld\n", rank, l); */
				eratosthenes[l/8] |= (1 << (l%8));
			}
		}

		/* If square of k is greater than the total gridsize then we don't need to carry on */
		if (pow(k,2) > gridsize)
			break;

	}
	if (rank+1 < size)
	{
		k = k*-1;
		MPI_Send(&k, 1, MPI_LONG, rank+1, TAG, MPI_COMM_WORLD);
	}

	c1=clock();

	k=gridpart;
	while((eratosthenes[k/8] & (1 << (k % 8))) != 0)
		k--;

	free(eratosthenes);
	MPI_Finalize(); /* Finalise MPI */

	printf ("Rank: %d, Done, largest prime %ld, %0.2f seconds\r\n", rank, k+start, (float)(c1 - c0)/CLOCKS_PER_SEC);
	return 0;
}

int initMPI(int *pargc, char *pargv[])
{
	int result;	

	/* Initialise MPI */
	result = MPI_Init(pargc, &pargv);
	if(result!=MPI_SUCCESS)
	{
		fprintf(stderr, "Error initialising MPI.");
		fflush(stderr);
		MPI_Abort(MPI_COMM_WORLD,result);
	}
	MPI_Comm_rank(MPI_COMM_WORLD, &rank); /* get rank */
	MPI_Comm_size(MPI_COMM_WORLD, &size); /* get size */
	return 0;
}
